# 加载数据集
import os.path

import numpy as np
from mindspore.dataset.vision import Inter

import preHandleData
import mindspore.dataset as ds
import mindspore.dataset.vision.c_transforms as c_vision


def create_dataset(data_dir, repeat=200, train_batch_size=16, augment=False, cross_val_ind=1, run_distribute=False):
    images = preHandleData._load_multipage_tiff(os.path.join(data_dir, 'volume.tif'))
    masks = preHandleData._load_multipage_tiff(os.path.join(data_dir, 'labels.tif'))

    train_indices, val_indices = preHandleData._get_val_train_indices(len(images), cross_val_ind)
    train_images = images[train_indices]
    train_masks = masks[train_images]

    train_images = np.repeat(train_masks, repeat, axis=0)
    train_masks = np.repeat(train_masks, repeat, axis=0)
    val_images = images[val_indices]
    val_masks = masks[val_indices]

    train_image_data = {'image': train_images}
    train_mask_data = {'mask': train_masks}
    valid_image_data = {'image': val_images}
    valid_mask_data = {'mask': val_masks}

    ds_train_images = ds.NumpySlicesDataset(data=train_image_data, sampler=None, shuffle=False)
    ds_train_masks = ds.NumpySlicesDataset(data=train_mask_data, sampler=None, shuffle=False)

    ds_valid_images = ds.NumpySlicesDataset(data=valid_image_data, sampler=None, shuffle=False)
    ds_valid_masks = ds.NumpySlicesDataset(data=valid_mask_data, sampler=None, shuffle=False)

    c_resize_op = c_vision.Resize(size=(388, 388), interpolation=Inter.BILINEAR)
    c_pad = c_vision.Pad(padding=92)
    c_rescale_image = c_vision.Rescale(1.0 / 127.5, -1)
    c_rescale_mask = c_vision.Rescale(1.0 / 255.0, 0)

    c_trans_normalize_img = [c_rescale_image, c_resize_op, c_pad]
    c_trans_normalize_mask = [c_rescale_mask, c_resize_op, c_pad]
    c_center_crop = c_vision.CenterCrop(size=388)

    train_image_ds = ds_train_images.map(input_columns='image', operations=c_trans_normalize_img)
    train_mask_ds = ds_train_masks.map(input_columns="mask", operations=c_trans_normalize_mask)
    train_ds = ds.zip((train_image_ds, train_mask_ds))
    train_ds = train_ds.project(columns=['image', 'mask'])

    if augment:
        augment_process = preHandleData.train_data_augmentation
        c_resize_op = c_vision.Resize(size=(572, 572), interpolation=Inter.BILINEAR)
        train_ds = train_ds.map(input_columns=['image', 'mask'], operations=c_resize_op)
        train_ds = train_ds.map(input_column='image', operations=c_resize_op)
        train_ds = train_ds.map(input_column='mask', operations=c_resize_op)

    train_ds = train_ds.map(input_columns='mask', operations=c_center_crop)
    post_process = preHandleData.data_post_process
    train_ds = train_ds.map(input_columns=['image', 'mask'], operations=post_process)
    train_ds = train_ds.shuffle(repeat * 24)
    train_ds = train_ds.batch(batch_size=train_batch_size, drop_remainder=True)

    valid_image_ds = ds_valid_images.map(input_columns='image', operations=c_trans_normalize_img)
    valid_mask_ds = ds_valid_masks.map(input_columns='mask', operations=c_trans_normalize_mask)
    valid_ds = ds.zip((valid_image_ds, valid_mask_ds))
    valid_ds = valid_ds.project(columns=['image', 'mask'])
    valid_ds = valid_ds.map(input_columns='mask', operations=c_center_crop)
    post_process = preHandleData.data_post_process
    valid_ds = valid_ds.map(input_columns=['image', 'mask'], operations=post_process)
    valid_ds = valid_ds.batch(batch_size=1, drop_remainder=True)

    return train_ds, valid_ds
